import sqlite3
from loguru import logger
from datetime import datetime

class AgentStatsDB:
    def __init__(self, db_path="memory_systems/agent_stats.db"):
        self.db_path = db_path
        self._create_table()

    def _create_table(self):
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS agent_stats (
                    agent_name TEXT PRIMARY KEY,
                    run_count INTEGER DEFAULT 0,
                    success_count INTEGER DEFAULT 0,
                    last_updated TEXT
                )
            """)
            conn.commit()
            logger.info(f"Database table 'agent_stats' created or already exists at {self.db_path}")
        except sqlite3.Error as e:
            logger.error(f"Error creating agent_stats table: {e}")
        finally:
            if conn:
                conn.close()

    def increment_run_count(self, agent_name: str, success: bool = True):
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            current_time = datetime.now().isoformat()

            cursor.execute("SELECT run_count, success_count FROM agent_stats WHERE agent_name = ?", (agent_name,))
            row = cursor.fetchone()

            if row:
                current_run_count, current_success_count = row
                new_run_count = current_run_count + 1
                new_success_count = current_success_count + (1 if success else 0)
                cursor.execute("""
                    UPDATE agent_stats SET
                        run_count = ?,
                        success_count = ?,
                        last_updated = ?
                    WHERE agent_name = ?
                """, (new_run_count, new_success_count, current_time, agent_name))
                logger.info(f"Updated stats for agent: {agent_name}")
            else:
                run_count = 1
                success_count = 1 if success else 0
                cursor.execute("""
                    INSERT INTO agent_stats (agent_name, run_count, success_count, last_updated)
                    VALUES (?, ?, ?, ?)
                """, (agent_name, run_count, success_count, current_time))
                logger.info(f"Inserted new stats for agent: {agent_name}")
            
            conn.commit()
        except sqlite3.Error as e:
            logger.error(f"Error incrementing run count for {agent_name}: {e}")
        finally:
            if conn:
                conn.close()

    def get_agent_stats(self, agent_name: str) -> dict | None:
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT run_count, success_count FROM agent_stats WHERE agent_name = ?", (agent_name,))
            row = cursor.fetchone()
            if row:
                return {"run_count": row[0], "success_count": row[1]}
            return None
        except sqlite3.Error as e:
            logger.error(f"Error retrieving stats for {agent_name}: {e}")
            return None
        finally:
            if conn:
                conn.close()

    def get_all_agent_stats(self) -> list[dict]:
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT agent_name, run_count, success_count FROM agent_stats")
            rows = cursor.fetchall()
            return [{"agent_name": row[0], "run_count": row[1], "success_count": row[2]} for row in rows]
        except sqlite3.Error as e:
            logger.error(f"Error retrieving all agent stats: {e}")
            return []
        finally:
            if conn:
                conn.close()

    def delete_agent_stats(self, agent_name: str):
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("DELETE FROM agent_stats WHERE agent_name = ?", (agent_name,))
            conn.commit()
            logger.info(f"Deleted stats for agent: {agent_name}")
        except sqlite3.Error as e:
            logger.error(f"Error deleting stats for {agent_name}: {e}")
        finally:
            if conn:
                conn.close()
